"""
流式处理模块 - 处理流式响应

此模块包含处理来自 LiteLLM 的流式响应并将其转换为 Anthropic 格式的函数。
"""

import json
import logging
import uuid
from models import MessagesRequest

logger = logging.getLogger(__name__)


async def handle_streaming(response_generator, original_request: MessagesRequest):
    """处理来自 LiteLLM 的流式响应并转换为 Anthropic 格式。"""
    try:
        # 发送 message_start 事件
        message_id = f"msg_{uuid.uuid4().hex[:24]}"  # 格式类似于 Anthropic 的 ID
        
        message_data = {
            'type': 'message_start',
            'message': {
                'id': message_id,
                'type': 'message',
                'role': 'assistant',
                'model': original_request.model,
                'content': [],
                'stop_reason': None,
                'stop_sequence': None,
                'usage': {
                    'input_tokens': 0,
                    'cache_creation_input_tokens': 0,
                    'cache_read_input_tokens': 0,
                    'output_tokens': 0
                }
            }
        }
        yield f"event: message_start\ndata: {json.dumps(message_data)}\n\n"
        
        # 第一个文本块的内容块索引
        yield f"event: content_block_start\ndata: {json.dumps({'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}})}\n\n"
        
        # 发送 ping 以保持连接活跃（Anthropic 这样做）
        yield f"event: ping\ndata: {json.dumps({'type': 'ping'})}\n\n"
        
        tool_index = None
        current_tool_call = None
        tool_content = ""
        accumulated_text = ""  # 跟踪累积的文本内容
        text_sent = False  # 跟踪是否已发送任何文本内容
        text_block_closed = False  # 跟踪文本块是否已关闭
        input_tokens = 0
        output_tokens = 0
        has_sent_stop_reason = False
        last_tool_index = 0
        
        # 处理每个块
        async for chunk in response_generator:
            try:

                
                # 检查这是否是带有使用数据的响应结束
                if hasattr(chunk, 'usage') and chunk.usage is not None:
                    if hasattr(chunk.usage, 'prompt_tokens'):
                        input_tokens = chunk.usage.prompt_tokens
                    if hasattr(chunk.usage, 'completion_tokens'):
                        output_tokens = chunk.usage.completion_tokens
                
                # 处理文本内容
                if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                    choice = chunk.choices[0]
                    
                    # 从选择中获取增量
                    if hasattr(choice, 'delta'):
                        delta = choice.delta
                    else:
                        # 如果没有增量，尝试获取消息
                        delta = getattr(choice, 'message', {})
                    
                    # 检查 finish_reason 以知道何时完成
                    finish_reason = getattr(choice, 'finish_reason', None)
                    
                    # 处理文本内容
                    delta_content = None
                    
                    # 处理不同格式的增量内容
                    if hasattr(delta, 'content'):
                        delta_content = delta.content
                    elif isinstance(delta, dict) and 'content' in delta:
                        delta_content = delta['content']
                    
                    # 累积文本内容
                    if delta_content is not None and delta_content != "":
                        accumulated_text += delta_content
                        
                        # 如果没有开始工具调用，总是发出文本增量
                        if tool_index is None and not text_block_closed:
                            text_sent = True
                            yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': delta_content}})}\n\n"
                    
                    # 处理工具调用
                    delta_tool_calls = None
                    
                    # 处理不同格式的工具调用
                    if hasattr(delta, 'tool_calls'):
                        delta_tool_calls = delta.tool_calls
                    elif isinstance(delta, dict) and 'tool_calls' in delta:
                        delta_tool_calls = delta['tool_calls']
                    
                    # 如果有工具调用，则处理
                    if delta_tool_calls:
                        # 我们看到的第一个工具调用 - 需要正确处理文本
                        if tool_index is None:
                            # 如果我们一直在流式传输文本，关闭该文本块
                            if text_sent and not text_block_closed:
                                text_block_closed = True
                                yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
                            # 如果我们已经累积了文本但未发送，现在需要发出它
                            # 这处理第一个增量同时包含文本和工具调用的情况
                            elif accumulated_text and not text_sent and not text_block_closed:
                                # 发送累积的文本
                                text_sent = True
                                yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': accumulated_text}})}\n\n"
                                # 关闭文本块
                                text_block_closed = True
                                yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
                            # 即使我们没有发送任何内容也要关闭文本块 - 模型有时会发出空文本块
                            elif not text_block_closed:
                                text_block_closed = True
                                yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
                                
                        # 如果还不是列表，则转换为列表
                        if not isinstance(delta_tool_calls, list):
                            delta_tool_calls = [delta_tool_calls]
                        
                        for tool_call in delta_tool_calls:
                            # 获取此工具调用的索引（用于多个工具）
                            current_index = None
                            if isinstance(tool_call, dict) and 'index' in tool_call:
                                current_index = tool_call['index']
                            elif hasattr(tool_call, 'index'):
                                current_index = tool_call.index
                            else:
                                current_index = 0
                            
                            # 检查这是新工具还是继续
                            if tool_index is None or current_index != tool_index:
                                # 新工具调用 - 创建新的 tool_use 块
                                tool_index = current_index
                                last_tool_index += 1
                                anthropic_tool_index = last_tool_index
                                
                                # 提取函数信息
                                if isinstance(tool_call, dict):
                                    function = tool_call.get('function', {})
                                    name = function.get('name', '') if isinstance(function, dict) else ""
                                    tool_id = tool_call.get('id', f"toolu_{uuid.uuid4().hex[:24]}")
                                else:
                                    function = getattr(tool_call, 'function', None)
                                    name = getattr(function, 'name', '') if function else ''
                                    tool_id = getattr(tool_call, 'id', f"toolu_{uuid.uuid4().hex[:24]}")
                                
                                # 开始新的 tool_use 块
                                yield f"event: content_block_start\ndata: {json.dumps({'type': 'content_block_start', 'index': anthropic_tool_index, 'content_block': {'type': 'tool_use', 'id': tool_id, 'name': name, 'input': {}}})}\n\n"
                                current_tool_call = tool_call
                                tool_content = ""
                            
                            # 提取函数参数
                            arguments = None
                            if isinstance(tool_call, dict) and 'function' in tool_call:
                                function = tool_call.get('function', {})
                                arguments = function.get('arguments', '') if isinstance(function, dict) else ''
                            elif hasattr(tool_call, 'function'):
                                function = getattr(tool_call, 'function', None)
                                arguments = getattr(function, 'arguments', '') if function else ''
                            
                            # 如果我们有参数，将它们作为增量发送
                            if arguments:
                                # 尝试检测参数是有效的 JSON 还是只是片段
                                try:
                                    # 如果已经是字典，使用它
                                    if isinstance(arguments, dict):
                                        args_json = json.dumps(arguments)
                                    else:
                                        # 否则，尝试解析它
                                        json.loads(arguments)
                                        args_json = arguments
                                except (json.JSONDecodeError, TypeError):
                                    # 如果是片段，将其视为字符串
                                    args_json = arguments
                                
                                # 添加到累积的工具内容
                                tool_content += args_json if isinstance(args_json, str) else ""
                                
                                # 发送更新
                                yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': anthropic_tool_index, 'delta': {'type': 'input_json_delta', 'partial_json': args_json}})}\n\n"
                    
                    # 处理 finish_reason - 结束流式响应
                    if finish_reason and not has_sent_stop_reason:
                        has_sent_stop_reason = True
                        
                        # 关闭任何打开的工具调用块
                        if tool_index is not None:
                            for i in range(1, last_tool_index + 1):
                                yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': i})}\n\n"
                        
                        # 如果我们累积了文本但从未发送或关闭文本块，现在就做
                        if not text_block_closed:
                            if accumulated_text and not text_sent:
                                # 发送累积的文本
                                yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': accumulated_text}})}\n\n"
                            # 关闭文本块
                            yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
                        
                        # 将 OpenAI finish_reason 映射到 Anthropic stop_reason
                        stop_reason = "end_turn"
                        if finish_reason == "length":
                            stop_reason = "max_tokens"
                        elif finish_reason == "tool_calls":
                            stop_reason = "tool_use"
                        elif finish_reason == "stop":
                            stop_reason = "end_turn"
                        
                        # 发送带有停止原因和使用情况的 message_delta
                        usage = {"output_tokens": output_tokens}
                        
                        yield f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': stop_reason, 'stop_sequence': None}, 'usage': usage})}\n\n"
                        
                        # 发送 message_stop 事件
                        yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
                        
                        # 发送最终的 [DONE] 标记以匹配 Anthropic 的行为
                        yield "data: [DONE]\n\n"
                        return
            except Exception as e:
                # 记录错误但继续处理其他块
                logger.error(f"处理块时出错：{str(e)}")
                continue
        
        # 如果我们没有得到结束原因，关闭任何打开的块
        if not has_sent_stop_reason:
            # 关闭任何打开的工具调用块
            if tool_index is not None:
                for i in range(1, last_tool_index + 1):
                    yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': i})}\n\n"
            
            # 关闭文本内容块
            yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
            
            # 发送带有使用情况的最终 message_delta
            usage = {"output_tokens": output_tokens}
            
            yield f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None}, 'usage': usage})}\n\n"
            
            # 发送 message_stop 事件
            yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
            
            # 发送最终的 [DONE] 标记以匹配 Anthropic 的行为
            yield "data: [DONE]\n\n"
    
    except Exception as e:
        import traceback
        error_traceback = traceback.format_exc()
        error_message = f"流式处理中出错：{str(e)}\n\n完整回溯：\n{error_traceback}"
        logger.error(error_message)
        
        # 发送错误 message_delta
        yield f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': 'error', 'stop_sequence': None}, 'usage': {'output_tokens': 0}})}\n\n"
        
        # 发送 message_stop 事件
        yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
        
        # 发送最终的 [DONE] 标记
        yield "data: [DONE]\n\n"